import time
import numpy as np
import win32gui, win32ui, win32con, win32api
from pynput.mouse import Controller
import cv2
import threading


def window_capture(hwnd,rate):

    hwndDC = win32gui.GetWindowDC(hwnd)
    # 根据窗口句柄获取窗口的设备上下文DC（Divice Context）
    mfcDC = win32ui.CreateDCFromHandle(hwndDC)
    # 根据窗口的DC获取mfcDC
    saveDC = mfcDC.CreateCompatibleDC()
    # mfcDC创建可兼容的DC    
    saveBitMap = win32ui.CreateBitmap()
    # 创建bigmap准备保存图片
    rctA = win32gui.GetWindowRect(hwnd)
    w = rctA[2] - rctA[0]
    h = rctA[3] - rctA[1]

    w = int(w*rate)
    h = int(h*rate)
    #获取图片大小
    # 截取从左上角（0，0）长宽为（w，h）的图片
    
    saveBitMap.CreateCompatibleBitmap(mfcDC, w, h)
    # 为bitmap开辟空间
    saveDC.SelectObject(saveBitMap)
    # 高度saveDC，将截图保存到saveBitmap中
    saveDC.BitBlt((0, 0), (w, h), mfcDC, (0, 0), win32con.SRCCOPY)
    
    signedIntsArray = saveBitMap.GetBitmapBits(True)
    img = np.frombuffer(signedIntsArray, dtype="uint8")
    img.shape = (h, w, 4)
    #bit图转mat图
    
    win32gui.DeleteObject(saveBitMap.GetHandle())
    mfcDC.DeleteDC()
    saveDC.DeleteDC()
    #释放内存
    return rctA,cv2.cvtColor(img, cv2.COLOR_RGBA2RGB) #转为RGB图返回





class shotThread(threading.Thread):
    def __init__(self,shot_time):
        threading.Thread.__init__(self)
        self.shot_time = shot_time
    def run(self):

        while True:
            if current_img is not None:
                
                # 图像名称
                img_name = int(time.time()).__str__() + ".png"
                # 保存图片到指定路径
                cv2.imwrite(img_name,current_img)

                time.sleep(self.shot_time)


 
def main(windows_name,rate,second):

    hwnd = win32gui.FindWindow(None, windows_name)

    # 定时截图线程可以不要这个逻辑,每两秒截图一张
    if second > 0:
        t = shotThread(second)
        t.setDaemon(True) 
        t.start()

    while True:
        # 截屏
        rctA,img = window_capture(hwnd,rate)

        # 游戏窗口的左上角的屏幕像素坐标
        window_x = rctA[0]
        window_y = rctA[1]

        global current_img
        current_img = img

        # 当前鼠标坐在的屏幕像素坐标
        mouseModal = Controller()
        x, y = mouseModal.position
        
        cv2.imshow("results", img)

        # 按英文q，结束录屏，或者ctrl + c
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        
    cv2.destroyAllWindows()






if __name__ == '__main__':

    # 存储最新一帧图像
    current_img = None

    # 要截屏的窗口名称(鼠标放到任务栏图标上可以显示名称)，如果是黑屏，要保证预览的应用在前台界面运行，不能最小化，则需要把被录屏/截图的程序重启:退出重新打开
    # windows_name = '微信'
    windows_name = '微信'
   
    # 屏幕缩放比例，要去windows的设置中查看，150% = 1.5，125% = 1.25
    rate = 1

    # 1秒截一张图,如果不需要截图,只预览画面，则最后一个参数传0
    shot_second = 0


    main(windows_name,rate,shot_second)